#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
多协议通信模块
支持HTTP/2、QUIC等多种协议的隐蔽通信
"""

import os
import sys
import json
import time
import ssl
import socket
import asyncio
import random
import base64
import threading
from typing import Dict, List, Optional, Any, Callable
from dataclasses import dataclass
from enum import Enum
import urllib.parse
import urllib.request

# 尝试导入HTTP/2和QUIC相关库
try:
    import h2.connection
    import h2.events
    import h2.config
    HTTP2_AVAILABLE = True
except ImportError:
    HTTP2_AVAILABLE = False
    print("警告: h2库未安装，HTTP/2功能不可用")

try:
    import aioquic
    from aioquic.asyncio import connect
    from aioquic.h3.connection import H3Connection
    from aioquic.quic.configuration import QuicConfiguration
    QUIC_AVAILABLE = True
except ImportError:
    QUIC_AVAILABLE = False
    print("警告: aioquic库未安装，QUIC功能不可用")

class ProtocolType(Enum):
    """协议类型枚举"""
    HTTP1 = "http1"
    HTTP2 = "http2"
    QUIC = "quic"
    WEBSOCKET = "websocket"
    DNS = "dns"
    ICMP = "icmp"

@dataclass
class CommunicationConfig:
    """通信配置"""
    protocol: ProtocolType
    server_host: str
    server_port: int
    encryption: bool = True
    compression: bool = False
    user_agent: str = ""
    headers: Dict[str, str] = None
    proxy: Optional[str] = None
    timeout: int = 30

class HTTP2Client:
    """HTTP/2客户端"""
    
    def __init__(self, config: CommunicationConfig):
        self.config = config
        self.connection = None
        self.socket = None
        
    def connect(self) -> bool:
        """建立HTTP/2连接"""
        try:
            if not HTTP2_AVAILABLE:
                return False
            
            # 创建SSL上下文
            ssl_context = ssl.create_default_context()
            ssl_context.set_alpn_protocols(['h2'])
            
            # 建立TCP连接
            self.socket = socket.create_connection(
                (self.config.server_host, self.config.server_port),
                timeout=self.config.timeout
            )
            
            # 包装SSL
            self.socket = ssl_context.wrap_socket(
                self.socket, 
                server_hostname=self.config.server_host
            )
            
            # 验证ALPN协商结果
            if self.socket.selected_alpn_protocol() != 'h2':
                print("服务器不支持HTTP/2")
                return False
            
            # 创建HTTP/2连接
            config = h2.config.H2Configuration(client_side=True)
            self.connection = h2.connection.H2Connection(config=config)
            self.connection.initiate_connection()
            
            # 发送连接前导
            self.socket.sendall(self.connection.data_to_send())
            
            return True
            
        except Exception as e:
            print(f"HTTP/2连接失败: {e}")
            return False
    
    def send_request(self, method: str, path: str, data: bytes = None) -> Optional[bytes]:
        """发送HTTP/2请求"""
        try:
            if not self.connection:
                return None
            
            # 构造请求头
            headers = [
                (':method', method),
                (':path', path),
                (':scheme', 'https'),
                (':authority', self.config.server_host),
                ('user-agent', self.config.user_agent or self._generate_user_agent()),
            ]
            
            # 添加自定义头
            if self.config.headers:
                for key, value in self.config.headers.items():
                    headers.append((key.lower(), value))
            
            # 发送请求
            stream_id = self.connection.get_next_available_stream_id()
            self.connection.send_headers(stream_id, headers)
            
            if data:
                self.connection.send_data(stream_id, data)
            
            self.connection.end_stream(stream_id)
            self.socket.sendall(self.connection.data_to_send())
            
            # 接收响应
            response_data = b''
            while True:
                data = self.socket.recv(65536)
                if not data:
                    break
                
                events = self.connection.receive_data(data)
                for event in events:
                    if isinstance(event, h2.events.DataReceived):
                        response_data += event.data
                    elif isinstance(event, h2.events.StreamEnded):
                        return response_data
                
                # 发送任何待发送的数据
                data_to_send = self.connection.data_to_send()
                if data_to_send:
                    self.socket.sendall(data_to_send)
            
            return response_data
            
        except Exception as e:
            print(f"HTTP/2请求失败: {e}")
            return None
    
    def _generate_user_agent(self) -> str:
        """生成随机User-Agent"""
        user_agents = [
            'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
            'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
            'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        ]
        return random.choice(user_agents)
    
    def close(self):
        """关闭连接"""
        if self.socket:
            self.socket.close()

class QUICClient:
    """QUIC客户端"""
    
    def __init__(self, config: CommunicationConfig):
        self.config = config
        self.connection = None
        
    async def connect(self) -> bool:
        """建立QUIC连接"""
        try:
            if not QUIC_AVAILABLE:
                return False
            
            # 创建QUIC配置
            configuration = QuicConfiguration(
                is_client=True,
                alpn_protocols=["h3"]
            )
            
            # 建立连接
            async with connect(
                self.config.server_host,
                self.config.server_port,
                configuration=configuration
            ) as protocol:
                self.connection = protocol
                return True
                
        except Exception as e:
            print(f"QUIC连接失败: {e}")
            return False
    
    async def send_request(self, method: str, path: str, data: bytes = None) -> Optional[bytes]:
        """发送QUIC请求"""
        try:
            if not self.connection:
                return None
            
            # 使用HTTP/3发送请求
            h3_conn = H3Connection(self.connection._quic)
            
            # 构造请求头
            headers = [
                (b":method", method.encode()),
                (b":path", path.encode()),
                (b":scheme", b"https"),
                (b":authority", self.config.server_host.encode()),
                (b"user-agent", (self.config.user_agent or self._generate_user_agent()).encode()),
            ]
            
            # 发送请求
            stream_id = h3_conn.get_next_available_stream_id()
            h3_conn.send_headers(stream_id, headers)
            
            if data:
                h3_conn.send_data(stream_id, data)
            
            # 接收响应
            response_data = b''
            # 这里需要实现具体的响应处理逻辑
            
            return response_data
            
        except Exception as e:
            print(f"QUIC请求失败: {e}")
            return None
    
    def _generate_user_agent(self) -> str:
        """生成随机User-Agent"""
        user_agents = [
            'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
            'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        ]
        return random.choice(user_agents)

class WebSocketClient:
    """WebSocket客户端"""
    
    def __init__(self, config: CommunicationConfig):
        self.config = config
        self.socket = None
        self.connected = False
        
    def connect(self) -> bool:
        """建立WebSocket连接"""
        try:
            import websocket
            
            # 构造WebSocket URL
            scheme = "wss" if self.config.encryption else "ws"
            url = f"{scheme}://{self.config.server_host}:{self.config.server_port}/"
            
            # 创建WebSocket连接
            self.socket = websocket.WebSocket()
            self.socket.connect(url, timeout=self.config.timeout)
            self.connected = True
            
            return True
            
        except Exception as e:
            print(f"WebSocket连接失败: {e}")
            return False
    
    def send_message(self, message: str) -> bool:
        """发送WebSocket消息"""
        try:
            if not self.connected:
                return False
            
            self.socket.send(message)
            return True
            
        except Exception as e:
            print(f"WebSocket发送失败: {e}")
            return False
    
    def receive_message(self) -> Optional[str]:
        """接收WebSocket消息"""
        try:
            if not self.connected:
                return None
            
            return self.socket.recv()
            
        except Exception as e:
            print(f"WebSocket接收失败: {e}")
            return None
    
    def close(self):
        """关闭WebSocket连接"""
        if self.socket:
            self.socket.close()
            self.connected = False

class DNSTunnelClient:
    """DNS隧道客户端"""
    
    def __init__(self, config: CommunicationConfig):
        self.config = config
        self.domain_suffix = ".example.com"  # 可配置的域名后缀
        
    def send_data(self, data: bytes) -> bool:
        """通过DNS查询发送数据"""
        try:
            # 将数据编码为Base32（DNS安全）
            import base64
            encoded_data = base64.b32encode(data).decode().lower()
            
            # 分割数据以适应DNS标签长度限制
            chunks = [encoded_data[i:i+50] for i in range(0, len(encoded_data), 50)]
            
            for i, chunk in enumerate(chunks):
                # 构造DNS查询域名
                query_domain = f"{i}.{chunk}{self.domain_suffix}"
                
                # 执行DNS查询
                import socket
                try:
                    socket.gethostbyname(query_domain)
                except socket.gaierror:
                    pass  # DNS查询失败是正常的，我们只是在传输数据
            
            return True
            
        except Exception as e:
            print(f"DNS隧道发送失败: {e}")
            return False

class MultiProtocolManager:
    """多协议管理器"""
    
    def __init__(self):
        self.clients = {}
        self.active_protocol = None
        self.fallback_protocols = [ProtocolType.HTTP1, ProtocolType.WEBSOCKET, ProtocolType.DNS]
        
    def create_client(self, config: CommunicationConfig):
        """创建协议客户端"""
        if config.protocol == ProtocolType.HTTP2:
            return HTTP2Client(config)
        elif config.protocol == ProtocolType.QUIC:
            return QUICClient(config)
        elif config.protocol == ProtocolType.WEBSOCKET:
            return WebSocketClient(config)
        elif config.protocol == ProtocolType.DNS:
            return DNSTunnelClient(config)
        else:
            # 默认HTTP/1.1
            return self._create_http1_client(config)
    
    def _create_http1_client(self, config: CommunicationConfig):
        """创建HTTP/1.1客户端"""
        class HTTP1Client:
            def __init__(self, config):
                self.config = config
            
            def connect(self):
                return True
            
            def send_request(self, method, path, data=None):
                try:
                    scheme = "https" if self.config.encryption else "http"
                    url = f"{scheme}://{self.config.server_host}:{self.config.server_port}{path}"
                    
                    req = urllib.request.Request(url, data=data, method=method)
                    
                    # 添加头部
                    if self.config.headers:
                        for key, value in self.config.headers.items():
                            req.add_header(key, value)
                    
                    with urllib.request.urlopen(req, timeout=self.config.timeout) as response:
                        return response.read()
                        
                except Exception as e:
                    print(f"HTTP/1.1请求失败: {e}")
                    return None
            
            def close(self):
                pass
        
        return HTTP1Client(config)
    
    def auto_select_protocol(self, server_host: str, server_port: int) -> ProtocolType:
        """自动选择最佳协议"""
        # 测试协议可用性
        protocols_to_test = [ProtocolType.HTTP2, ProtocolType.QUIC, ProtocolType.WEBSOCKET, ProtocolType.HTTP1]
        
        for protocol in protocols_to_test:
            if self._test_protocol_availability(protocol, server_host, server_port):
                return protocol
        
        # 如果都不可用，返回DNS隧道作为最后手段
        return ProtocolType.DNS
    
    def _test_protocol_availability(self, protocol: ProtocolType, host: str, port: int) -> bool:
        """测试协议可用性"""
        try:
            config = CommunicationConfig(
                protocol=protocol,
                server_host=host,
                server_port=port,
                timeout=5
            )
            
            client = self.create_client(config)
            
            if protocol == ProtocolType.QUIC:
                # QUIC需要异步测试
                return asyncio.run(self._test_quic_async(client))
            else:
                return client.connect()
                
        except Exception:
            return False
    
    async def _test_quic_async(self, client) -> bool:
        """异步测试QUIC连接"""
        try:
            return await client.connect()
        except:
            return False
    
    def send_with_fallback(self, configs: List[CommunicationConfig], 
                          method: str, path: str, data: bytes = None) -> Optional[bytes]:
        """使用回退机制发送请求"""
        for config in configs:
            try:
                client = self.create_client(config)
                
                if config.protocol == ProtocolType.QUIC:
                    # QUIC需要异步处理
                    return asyncio.run(self._send_quic_async(client, method, path, data))
                else:
                    if client.connect():
                        response = client.send_request(method, path, data)
                        client.close()
                        if response:
                            return response
                            
            except Exception as e:
                print(f"协议 {config.protocol.value} 失败: {e}")
                continue
        
        return None
    
    async def _send_quic_async(self, client, method: str, path: str, data: bytes = None) -> Optional[bytes]:
        """异步发送QUIC请求"""
        try:
            if await client.connect():
                return await client.send_request(method, path, data)
        except Exception as e:
            print(f"QUIC异步请求失败: {e}")
        return None
    
    def create_protocol_chain(self, primary_protocol: ProtocolType, 
                            server_host: str, server_port: int) -> List[CommunicationConfig]:
        """创建协议链（主协议+回退协议）"""
        configs = []
        
        # 主协议
        primary_config = CommunicationConfig(
            protocol=primary_protocol,
            server_host=server_host,
            server_port=server_port,
            encryption=True
        )
        configs.append(primary_config)
        
        # 回退协议
        for fallback_protocol in self.fallback_protocols:
            if fallback_protocol != primary_protocol:
                fallback_config = CommunicationConfig(
                    protocol=fallback_protocol,
                    server_host=server_host,
                    server_port=server_port,
                    encryption=True
                )
                configs.append(fallback_config)
        
        return configs

# 使用示例
if __name__ == "__main__":
    # 创建多协议管理器
    manager = MultiProtocolManager()
    
    # 自动选择协议
    best_protocol = manager.auto_select_protocol("example.com", 443)
    print(f"推荐协议: {best_protocol.value}")
    
    # 创建协议链
    protocol_chain = manager.create_protocol_chain(
        ProtocolType.HTTP2, "example.com", 443
    )
    
    # 发送请求（带回退）
    test_data = b"Hello, World!"
    response = manager.send_with_fallback(
        protocol_chain, "POST", "/api/test", test_data
    )
    
    if response:
        print(f"响应: {response[:100]}...")
    else:
        print("所有协议都失败了")